#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
无人机飞联网传输协议仿真系统 - 主入口
"""

import os
import argparse
import sys
import time
from datetime import datetime
import json
import matplotlib.pyplot as plt
import numpy as np

# 确保导入路径正确
current_dir = os.path.dirname(os.path.abspath(__file__))
if current_dir not in sys.path:
    sys.path.append(current_dir)

# 确保可以导入模块
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from src.simulation import UAVCommunicationSimulation
from src.visualizer import UAVSimulationVisualizer

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='UAV通信协议仿真与可视化')
    
    parser.add_argument('--num_uavs', type=int, default=20,
                       help='无人机数量 (默认: 20)')
    parser.add_argument('--num_rounds', type=int, default=1000,
                       help='仿真轮次 (默认: 1000)')
    parser.add_argument('--emergency_rate', type=float, default=0.1,
                       help='紧急消息比例 (默认: 0.1)')
    parser.add_argument('--deadline_rate', type=float, default=0.3,
                       help='带截止时间消息比例 (默认: 0.3)')
    parser.add_argument('--seed', type=int, default=None,
                       help='随机种子 (默认: None)')
    parser.add_argument('--output_dir', type=str, default='./results',
                       help='输出目录 (默认: ./results)')
    parser.add_argument('--protocols', type=str, nargs='+',
                       default=['traditional', 'zigbee', 'wifi_direct', 'lora', 'fanet', 'dcdp_protocol'],
                       help='要测试的协议列表 (默认: 所有协议)')
    parser.add_argument('--mobility_model', type=str, default='random_waypoint',
                       choices=['random_waypoint', 'gauss_markov'],
                       help='无人机移动模型 (默认: random_waypoint)')
    parser.add_argument('--area_size', type=int, default=1000,
                       help='仿真区域大小(米) (默认: 1000)')
    parser.add_argument('--mobility_update_interval', type=int, default=10,
                       help='无人机位置更新间隔(轮次) (默认: 10)')
    parser.add_argument('--track_trajectories', action='store_true',
                       help='是否记录无人机移动轨迹')
    parser.add_argument('--show_dcdp_delay_analysis', action='store_true',
                       help='显示DCDP三类时延详细分析')
    parser.add_argument('--separate_metrics', action='store_true',
                       help='绘制四个单独的性能指标图表（不包含lora协议）')
    parser.add_argument('--exclude_lora', action='store_true',
                       help='从所有图表中排除lora协议')
    
    return parser.parse_args()

def print_simulation_banner(args):
    """打印仿真开始信息"""
    print("="*80)
    print(f"UAV通信协议仿真开始 - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print("-"*80)
    print(f"无人机数量: {args.num_uavs}")
    print(f"仿真轮次: {args.num_rounds}")
    print(f"紧急消息比例: {args.emergency_rate}")
    print(f"截止时间消息比例: {args.deadline_rate}")
    print(f"随机种子: {args.seed}")
    print(f"仿真区域大小: {args.area_size}m x {args.area_size}m")
    print(f"无人机移动模型: {args.mobility_model}")
    print(f"位置更新间隔: 每{args.mobility_update_interval}轮")
    print(f"是否记录移动轨迹: {'是' if args.track_trajectories else '否'}")
    print(f"测试协议: {', '.join(args.protocols)}")
    print("="*80)
    print()

def print_simulation_results(results):
    """打印仿真结果概要"""
    print("\n" + "="*80)
    print("Simulation Results Summary:")
    print("-"*80)
    
    # 提取需要显示的统计信息
    stats = results['statistics']
    protocols = list(stats.keys())
    
    # 计算成功率
    success_rates = {p: (stats[p]['successful_messages'] / stats[p]['total_messages'] * 100) 
                    for p in protocols}
    
    # 计算紧急消息处理率
    emergency_rates = {p: (stats[p]['emergency_messages_handled'] / stats[p]['emergency_messages'] * 100 
                         if stats[p]['emergency_messages'] > 0 else 0) 
                      for p in protocols}
    
    # 计算截止时间满足率
    deadline_rates = {p: (stats[p]['deadline_met'] / stats[p]['deadline_messages'] * 100 
                        if stats[p]['deadline_messages'] > 0 else 0) 
                     for p in protocols}
    
    # 按成功率排序协议
    sorted_protocols = sorted(protocols, key=lambda p: success_rates[p], reverse=True)
    
    # 打印表头
    print(f"{'Protocol':<15} {'Success Rate':<12} {'Emergency Rate':<15} {'Deadline Met':<15} {'Avg Delay(ms)':<15}")
    print("-"*80)
    
    # 打印每个协议的结果
    for protocol in sorted_protocols:
        print(f"{protocol:<15} {success_rates[protocol]:>10.2f}% {emergency_rates[protocol]:>13.2f}% "
              f"{deadline_rates[protocol]:>13.2f}% {stats[protocol]['avg_delay']:>13.2f}")
    
    print("-"*80)
    print(f"Simulation Duration: {results['simulation_params']['duration']:.2f} seconds")
    
    # 如果有DCDP协议，打印三类时延分析
    if 'dcdp_protocol' in stats and 'avg_propagation_delay' in stats['dcdp_protocol']:
        dcdp_stats = stats['dcdp_protocol']
        print("\nDCDP Protocol Delay Components Analysis:")
        print(f"Propagation Delay: {dcdp_stats['avg_propagation_delay']:.2f} ms")
        print(f"Processing Delay:  {dcdp_stats['avg_processing_delay']:.2f} ms")
        print(f"Congestion Delay:  {dcdp_stats['avg_congestion_delay']:.2f} ms")
        total_delay = (dcdp_stats['avg_propagation_delay'] + 
                      dcdp_stats['avg_processing_delay'] + 
                      dcdp_stats['avg_congestion_delay'])
        print(f"Total Delay:       {total_delay:.2f} ms")
        print(f"Delay Proportion:")
        print(f"  - Propagation:   {(dcdp_stats['avg_propagation_delay']/total_delay*100):.2f}%")
        print(f"  - Processing:    {(dcdp_stats['avg_processing_delay']/total_delay*100):.2f}%")
        print(f"  - Congestion:    {(dcdp_stats['avg_congestion_delay']/total_delay*100):.2f}%")
    
    print("="*80)
    print()

def track_uav_mobility(simulator, num_rounds, update_interval):
    """记录无人机移动轨迹
    
    Args:
        simulator: 仿真器实例
        num_rounds: 仿真总轮次
        update_interval: 位置更新间隔
        
    Returns:
        轨迹数据和距离数据
    """
    # 初始化轨迹跟踪
    trajectories = [[] for _ in range(simulator.num_uavs)]
    avg_distances = []
    
    # 记录初始位置
    for i in range(simulator.num_uavs):
        # 使用列表而不是np.ndarray
        pos = simulator.uav_positions[i].tolist()
        trajectories[i].append(pos)
    
    # 计算初始平均距离
    total_dist = 0
    count = 0
    for i in range(simulator.num_uavs):
        for j in range(i+1, simulator.num_uavs):
            total_dist += simulator.calculate_distance(i, j)
            count += 1
    avg_distances.append(total_dist / max(1, count))
    
    # 模拟无人机移动并记录轨迹
    for round_num in range(0, num_rounds, update_interval):
        # 更新位置
        simulator.update_uav_positions(step_time=1.0)
        
        # 记录位置
        for i in range(simulator.num_uavs):
            # 使用列表而不是np.ndarray
            pos = simulator.uav_positions[i].tolist()
            trajectories[i].append(pos)
        
        # 计算平均距离
        total_dist = 0
        count = 0
        for i in range(simulator.num_uavs):
            for j in range(i+1, simulator.num_uavs):
                total_dist += simulator.calculate_distance(i, j)
                count += 1
        avg_distances.append(total_dist / max(1, count))
    
    return trajectories, avg_distances

def main():
    """主函数"""
    # 解析命令行参数
    args = parse_args()
    
    # 打印仿真开始信息
    print_simulation_banner(args)
    
    # 确保输出目录存在
    os.makedirs(args.output_dir, exist_ok=True)
    os.makedirs(os.path.join(args.output_dir, 'figures'), exist_ok=True)
    
    # 如果需要排除lora协议
    if args.exclude_lora and 'lora' in args.protocols:
        args.protocols.remove('lora')
        print("已从协议列表中排除lora协议")
    
    # 创建仿真器
    simulator = UAVCommunicationSimulation(
        num_uavs=args.num_uavs, 
        seed=args.seed,
        area_size=args.area_size,
        mobility_model=args.mobility_model
    )
    
    # 如果需要跟踪轨迹，先记录轨迹
    trajectories = None
    distances = None
    if args.track_trajectories:
        print("记录无人机移动轨迹...")
        trajectories, distances = track_uav_mobility(
            simulator, args.num_rounds, args.mobility_update_interval)
    
    # 运行仿真
    print("正在运行仿真...")
    results = simulator.run_simulation(
        num_rounds=args.num_rounds,
        emergency_rate=args.emergency_rate,
        deadline_rate=args.deadline_rate,
        mobility_update_interval=args.mobility_update_interval
    )
    
    # 如果记录了轨迹，添加到结果中
    if args.track_trajectories and trajectories and distances:
        results['uav_trajectories'] = trajectories
        results['uav_distances'] = distances
    
    # 保存结果
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    results_filename = os.path.join(args.output_dir, f"sim_results_{timestamp}.json")
    simulator.save_results(results, results_filename)
    
    # 打印结果概要
    print_simulation_results(results)
    
    # 可视化结果
    print("正在生成可视化结果...")
    figures_dir = os.path.join(args.output_dir, 'figures')
    visualizer = UAVSimulationVisualizer(results, figures_dir)
    
    # 如果指定生成单独的指标图表
    if args.separate_metrics:
        visualizer.plot_separate_metrics()
    else:
        visualizer.plot_all()
    
    print(f"可视化结果已保存至: {figures_dir}")
    print("\n仿真完成！")

if __name__ == "__main__":
    main() 